# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment3/question1/mediums.py
import abc
import typing

import numpy
import scipy.constants


class Medium(abc.ABC):
    @abc.abstractmethod
    def transfer_matrix(self,
                        circular_frequency: float,
                        incident_angle: float) -> numpy.ndarray:
        raise NotImplementedError("The method should be implemented by its subclasses.")

    def concatenate(self, other: 'Medium') -> 'Medium':
        class ConcatenatedMedium(Medium):
            def __init__(self, left: Medium, right: Medium):
                self._left = left
                self._right = right

            def transfer_matrix(self,
                                circular_frequency: float,
                                incident_angle: float) -> numpy.ndarray:
                left = self._left.transfer_matrix(
                    circular_frequency, incident_angle)
                right = self._right.transfer_matrix(
                    circular_frequency, incident_angle)
                return numpy.matmul(left, right)

        return ConcatenatedMedium(other, self)

    @staticmethod
    def create(width: float, refractive_index: complex) -> 'Medium':
        class CreatedMedium(Medium):
            def __init__(self, h: float, n: complex):
                self._h = h
                self._n = n

            def transfer_matrix(self,
                                circular_frequency: float,
                                incident_angle: float) -> numpy.ndarray:
                delta = self._n
                delta *= self._h
                delta *= numpy.cos(incident_angle)
                delta *= circular_frequency
                delta /= scipy.constants.c

                c = numpy.cos(delta)
                s = numpy.sin(delta)
                return numpy.array([
                    [c, s * self._n * 1j],
                    [s / self._n * 1j, c]
                ])

        return CreatedMedium(width, refractive_index)

    @staticmethod
    def join(mediums: typing.Iterable['Medium']) -> typing.Union['Medium', None]:
        result: Medium | None = None
        for medium in mediums:
            if not result:
                result = medium
            else:
                result = result.concatenate(medium)
        return result
